// Reference:
// https://github.com/TheAlgorithms/Go/blob/master/structure/avl/avl.go
// avl is self-balancing tree, i.e. for all node in a tree, height difference
// between its left and right child will not exceed 1.
package avl

type Node struct {
	Key         int
	Height      int
	Left, Right *Node
}

// Creates a new Tree
func NewTree() *Node {
	return nil
}

// Returns node with given key
func Get(root *Node, key int) *Node {
	if root == nil {
		return root
	}

	if root.Key == key {
		return root
	} else if root.Key < key {
		root = root.Right
	} else {
		root = root.Left
	}

	return Get(root, key)
}

func Insert(root **Node, key int) {
	if *root == nil {
		*root = &Node{
			Key:    key,
			Height: 1,
		}
		return
	}

	if (*root).Key < key {
		Insert(&(*root).Right, key)
	} else if (*root).Key > key {
		Insert(&(*root).Left, key)
	}

	(*root).Height = height(*root)
	bFactor := balanceFactor(*root)

	switch bFactor {
	case 2:
		lFactor := balanceFactor((*root).Left)
		switch lFactor {
		case 1:
			llRotation(root)
		case -1:
			lrRotation(root)
		}
	case -2:
		rFactor := balanceFactor((*root).Right)
		switch rFactor {
		case 1:
			rlRotation(root)
		case -1:
			rrRotation(root)
		}
	}
}

func Delete(root **Node, key int) {
	if root == nil {
		return
	}

	if (*root).Key < key {
		Delete(&(*root).Right, key)
	} else if (*root).Key > key {
		Delete(&(*root).Left, key)
	} else {
		if (*root).Left == nil && (*root).Right == nil {
			*root = nil
		} else if (*root).Left == nil {
			*root = (*root).Right
		} else if (*root).Right == nil {
			*root = (*root).Left
		} else {
			minVal := min((*root).Right)
			(*root).Key = minVal
			Delete(root, minVal)
		}
		return
	}

	(*root).Height = height(*root)
	bFactor := balanceFactor(*root)

	switch bFactor {
	case 2:
		switch balanceFactor((*root).Left) {
		case 1:
			llRotation(root)
		case -1:
			lrRotation(root)
		case 0:
			llRotation(root)
		}
	case -2:
		switch balanceFactor((*root).Right) {
		case 1:
			rlRotation(root)
		case -1:
			rrRotation(root)
		case 0:
			rrRotation(root)
		}
	}
}

func llRotation(root **Node) {
	b := (*root).Left
	br := b.Right
	b.Right = *root
	(*root).Left = br

	(*root).Height = height(*root)
	b.Height = height(b)

	*root = b
}

func lrRotation(root **Node) {
	c := (*root).Left.Right
	cl := c.Left
	cr := c.Right

	c.Left = (*root).Left
	c.Right = (*root)
	c.Left.Right = cl

	(*root).Left = cr
	(*root).Height = height(*root)
	c.Left.Height = height(c.Left)
	c.Height = height(c)

	*root = c
}

func rrRotation(root **Node) {
	b := (*root).Right
	bl := b.Left
	b.Left = *root

	(*root).Right = bl
	(*root).Height = height(*root)
	b.Height = height(b)
	*root = b
}

func rlRotation(root **Node) {
	c := (*root).Right.Left
	cl := c.Left
	cr := c.Right

	c.Right = (*root).Right
	c.Right.Left = cr
	c.Left = *root
	(*root).Right = cl

	(*root).Height = height(*root)
	c.Right.Height = height(c.Right)
	c.Height = height(c)
	*root = c
}

func balanceFactor(root *Node) int {
	var left, right int
	if root.Left != nil {
		left = root.Left.Height
	}
	if root.Right != nil {
		right = root.Right.Height
	}

	return left - right
}

func height(root *Node) int {
	var calculate func(root *Node, depth int) int

	calculate = func(root *Node, depth int) int {
		if root == nil {
			return depth
		}

		return Max(calculate(root.Left, depth+1), calculate(root.Right, depth+1))
	}

	return calculate(root, 0)
}

func Max(a, b int) int {
	if a >= b {
		return a
	}
	return b
}

func min(root *Node) int {
	if root.Left == nil {
		return root.Key
	}
	return min(root.Left)
}
